import OpenAI from 'openai';
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";

// 使用硅基流动的API配置
const client = new OpenAI({
  apiKey: "<替换为你的API Key>",
  baseURL: "https://api.siliconflow.cn/v1"
});

// 构建单个MCP服务器连接管理类
class MCPClient {
  // 构造函数
  constructor() {
    this.mcp = new Client({ name: "mcp-client-simple", version: "1.0.0" });
    this.transport = null;
    this.tools = [];
  }

  // 请求大模型
  async fetchLLM(messages){
    return await client.chat.completions.create({
      model: "Qwen/Qwen3-8B",
      max_tokens: 1000,
      enable_thinking: false,
      messages,
      tools: this.tools.map(tool => ({
        type: 'function',
        function: {
          name: tool.name,
          description: tool.description,
          parameters: tool.input_schema
        }
      })),
    });
  }
  
  // 连接到MCP服务器
  async connectToServer(serverScriptPath) {
    /**
     * 连接到MCP服务器
     * @param {string} serverScriptPath - 服务器脚本路径
     */
    try {
      // 确定脚本类型和适当的命令
      const isJs = serverScriptPath.endsWith(".js");
      const isPy = serverScriptPath.endsWith(".py");
      if (!isJs && !isPy) {
        throw new Error("Server script must be a .js or .py file");
      }
      const command = isPy
        ? process.platform === "win32"
          ? "python"
          : "python3"
        : process.execPath;

      // 初始化传输并连接到服务器
      this.transport = new StdioClientTransport({
        command,
        args: [serverScriptPath],
      });
      await this.mcp.connect(this.transport);

      // 列出可用工具
      const toolsResult = await this.mcp.listTools();
      // 获取服务器版本信息
      const serverVersion = await this.mcp.getServerVersion()
      this.tools = toolsResult.tools.map((tool) => {
        return {
          name: tool.name,
          description: tool.description,
          input_schema: tool.inputSchema,
        };
      });
      console.log(
        `已连接到MCP服务器:${serverVersion.name}，发现如下可用工具：`,
        this.tools.map(({ name }) => name),
      );
    } catch (e) {
      console.log("Failed to connect to MCP server: ", e);
      throw e;
    }
  }
  
  // 处理问题
  async processQuery(query) {
    /**
     * 使用LLM和可用工具处理问题
     * @param {string} query - 用户的输入问题
     * @returns {string} 处理后的响应
     */
    const messages = [
      {
        role: "user",
        content: query,
      },
    ];

    // 首次调用大模型，根据“用户的提问内容”结合“可用工具函数列表信息”进行【意图识别】并【给出工具函数调用决策】）
    const response = await this.fetchLLM(messages);

    // 处理响应和工具调用
    const message = response.choices[0].message;
    const finalText = [];
    const toolResults = [];

    // 如果有文本内容
    if (message.content) {
      finalText.push(message.content);
    }

    // 如果有工具调用
    if (message.tool_calls) {
      // 遍历所有的工具调用
      console.log(`\n🔧 开始工具调用： `)  
      for (const toolCall of message.tool_calls) {
        const toolName = toolCall.function.name;
        const toolInput = JSON.parse(toolCall.function.arguments);

        console.log(`调用工具"${toolName}",传入${JSON.stringify(toolInput)}`);

        try {
          // 调用工具
          const result = await this.mcp.callTool({
            name: toolName,
            arguments: toolInput,
          });

          const toolOutput = result.content[0]?.text || JSON.stringify(result.content);
          console.log(`获取结果:`, toolOutput);

          // 将工具调用结果添加到结果列表
          toolResults.push({
            tool_call_id: toolCall.id,
            role: "tool",
            content: toolOutput,
          });
        } catch (error) {
          console.error(`Tool execution failed:`, error);
          // 如果工具调用失败，将错误信息添加到结果列表
          toolResults.push({
            tool_call_id: toolCall.id,
            role: "tool",
            content: `Error: ${error.message}`,
          });
        }
      }
      console.log(`🔧 结束工具调用。\n`)  
    }

    // 如果有工具调用，需要再次调用API获取最终结果
    if (toolResults.length > 0) {

      // 将工具调用结果添加到消息列表
      messages.push({
        role: "assistant",
        content: message.content,
        tool_calls: message.tool_calls,
      });

      // 添加工具调用结果
      messages.push(...toolResults);
      
      // 再次调用大模型，整合【用户原始问题】、【函数调用决策】及【函数调用结果】给出最终结果。
      const finalResponse = await this.fetchLLM(messages);

      return finalResponse.choices[0].message.content || "";
    }

    return finalText.join("\n");
  }
  
  // 断开连接
  async disconnect() {
    if (this.transport) {
      await this.transport.close();
    }
  }
}

// 主程序
async function main() {
  const mcpClient = new MCPClient();

  try {
    // 连接到MCP服务器
    await mcpClient.connectToServer("./server/server-multi-tools.js");
    // 测试问题
    const prompt = "今天几号？9加7等于多少？请告诉我当前IP地址和所在城市？";
    // 开始提问
    console.log(`\n问题: ${prompt}`);
    const result = await mcpClient.processQuery(prompt);
    console.log(`回答: ${result}`);

  } catch (error) {
    console.error("Error:", error);
  } finally {
    // 断开连接
    await mcpClient.disconnect();
    console.log(`断开连接。\n`);
  }
}

// 运行主程序
main().catch(console.error);